from math import log
import torch
from scipy.optimize import linear_sum_assignment
from torch import nn

from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou


def build_matcher(args, loss_type='ce_loss'):
    '''
        匈牙利算法做匹配
    '''
    #   N, 100, xxx
    return HungarianMatcherHOI(cost_obj_class=args.set_cost_obj_class, cost_verb_class=args.set_cost_verb_class,
                               cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou, 
                               cost_matching=args.set_cost_matching, use_matching=args.use_matching,
                               loss_type=loss_type)

class HungarianMatcherHOI(nn.Module):
    def __init__(self, cost_obj_class: float = 1, cost_verb_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1,
                 cost_matching:float = 1, use_matching: bool = False, loss_type='ce_loss'):
        super().__init__()
        self.obj_alpha = 0.25
        self.obj_gamma = 2
        self.verb_alpha = 0.5
        self.verb_gamma = 2
        self.neg_inds_beta = 0
        self.hard_neg_inds_beta = 3
        
        self.loss_type = loss_type
        
        self.cost_obj_class = cost_obj_class
        self.cost_verb_class = cost_verb_class
        self.cost_matching = cost_matching
        self.cost_giou = cost_giou
        self.cost_bbox = cost_bbox
        self.use_matching = use_matching
        assert cost_obj_class != 0 or cost_verb_class != 0 or cost_bbox != 0 or cost_giou != 0, 'all costs cant be 0'

    @torch.no_grad()
    def forward(self, outputs, targets):
        C = self.get_C(outputs, targets)       
        sizes = [len(v['obj_labels']) for v in targets]
        #   list长度为N，i是pred，j是GT,    寻找loss最小的(不会重复)
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
        indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
        return indices


    @torch.no_grad()
    def get_C(self, outputs, targets):
        bs, num_queries = outputs['pred_obj_logits'].shape[:2]         
        C = self.get_cost_obj_class(outputs['pred_obj_logits'], targets,  self.obj_alpha, self.obj_gamma, self.neg_inds_beta, self.hard_neg_inds_beta) + \
            self.get_cost_bbox(outputs['pred_sub_boxes'], outputs['pred_obj_boxes'], targets) + \
            self.get_cost_giou(outputs['pred_sub_boxes'], outputs['pred_obj_boxes'], targets)  + \
            self.get_cost_verb_class(outputs['pred_verb_logits'], targets,  self.verb_alpha, self.verb_gamma, self.neg_inds_beta, self.hard_neg_inds_beta) 
        #   主力匹配还是get_cost_bbox,优先召回率，然后精确度
        C = C.view(bs, num_queries, -1).cpu()
        return C    

    def get_cost_obj_class(self, pred_obj_logits, targets,  
                           alpha=0.25, gamma=2, neg_inds_beta=0, hard_neg_inds_beta=3):
        #   100N, 29
        out_obj_prob = pred_obj_logits.flatten(0, 1).softmax(-1)
        #   tgtN,
        tgt_obj_labels = torch.cat([v['obj_labels'] for v in targets])

        
        if self.loss_type=='ce_loss':
            cost_obj_class = -out_obj_prob[:, tgt_obj_labels]#1-out_obj_prob[:, tgt_obj_labels]
        elif self.loss_type=='focal_loss':
            neg_cost = (1 - alpha) * (out_obj_prob ** gamma) * (-(1 - out_obj_prob + 1e-8).log())
            pos_cost = alpha * ((1 - out_obj_prob) ** gamma) * (-(out_obj_prob + 1e-8).log())
            cost_obj_class = pos_cost[:, tgt_obj_labels] - neg_cost[:, tgt_obj_labels]            

        C = self.cost_obj_class * cost_obj_class
        return C

    def get_cost_verb_class(self, pred_verb_logits, targets, alpha=0.5, gamma=2, neg_inds_beta=0, hard_neg_inds_beta=3):
        out_verb_prob = pred_verb_logits.flatten(0, 1).sigmoid()
        tgt_verb_labels = torch.cat([v['verb_labels'] for v in targets])
        tgt_verb_labels = tgt_verb_labels.permute(1, 0)
        
        if self.loss_type=='ce_loss':  
            pos = out_verb_prob.matmul(tgt_verb_labels)
            pos_sum = tgt_verb_labels.sum(dim=0, keepdim=True) + 1e-4
            neg = (1 - out_verb_prob).matmul(1 - tgt_verb_labels)
            neg_num = ((1 - tgt_verb_labels).sum(dim=0, keepdim=True) + 1e-4)
            
            cost_verb_class = -(pos / pos_sum + neg / neg_num) / 2
            
        elif self.loss_type=='focal_loss':
            #out_verb_prob = torch.clamp(out_verb_prob, 1e-12) 
            pos_pred = alpha * ((1 - out_verb_prob) ** gamma) * (-(out_verb_prob + 1e-8).log())
            pos_inds = tgt_verb_labels.eq(1).float()
            pos_target = pos_inds 
            pos_verb_cost_class = pos_pred.matmul(pos_target)
            pos_tgt_verb_labels_sum = pos_inds.sum(dim=0, keepdim=True) + 1e-4
            
            neg_pred= (1 - alpha) * (out_verb_prob ** gamma) * (-(1 - out_verb_prob + 1e-8).log())
            neg_inds = (tgt_verb_labels.lt(1) & tgt_verb_labels.gt(-1)).float()
            hard_neg_inds = (tgt_verb_labels.eq(-1)).float()
            neg_weights = torch.pow(1 - tgt_verb_labels, neg_inds_beta) * neg_inds + \
                        torch.pow(1 - tgt_verb_labels, hard_neg_inds_beta) * hard_neg_inds #   2
            neg_target = neg_weights 
            neg_verb_cost_class = neg_pred.matmul(neg_target)
            neg_tgt_verb_labels_sum = (neg_inds + hard_neg_inds).sum(dim=0, keepdim=True) + 1e-4
            cost_verb_class = (pos_verb_cost_class  + neg_verb_cost_class) / pos_tgt_verb_labels_sum
                    
        #   100N, tgtN
        C = self.cost_verb_class * cost_verb_class
        return C

    def get_cost_bbox(self, pred_sub_boxes, pred_obj_boxes, targets):
        #   100N, 4
        out_sub_bbox = pred_sub_boxes.flatten(0, 1)
        tgt_sub_boxes = torch.cat([v['sub_boxes'] for v in targets])
        cost_sub_bbox = torch.cdist(out_sub_bbox, tgt_sub_boxes, p=1)
                   
        out_obj_bbox = pred_obj_boxes.flatten(0, 1)
        tgt_obj_boxes = torch.cat([v['obj_boxes'] for v in targets])
        cost_obj_bbox = torch.cdist(out_obj_bbox, tgt_obj_boxes, p=1) * \
                        (tgt_obj_boxes != 0).any(dim=1).unsqueeze(0)#   tgt_cls = 0不计算cost
        
        if cost_sub_bbox.shape[1] == 0:#图片上没有物体
            cost_bbox = cost_sub_bbox
        else:
            cost_bbox = torch.stack((cost_sub_bbox, cost_obj_bbox)).max(dim=0)[0]

        #   100N, tgtN
        C = self.cost_bbox * cost_bbox 
        return C


    def get_cost_giou(self, pred_sub_boxes, pred_obj_boxes, targets):
        #   100N, 4
        out_sub_bbox = pred_sub_boxes.flatten(0, 1)
        tgt_sub_boxes = torch.cat([v['sub_boxes'] for v in targets])
        bbox1 = box_cxcywh_to_xyxy(out_sub_bbox)
        bbox2 = box_cxcywh_to_xyxy(tgt_sub_boxes)
        cost_sub_giou = -generalized_box_iou(bbox1, bbox2)
        
        
        out_obj_bbox = pred_obj_boxes.flatten(0, 1)
        tgt_obj_boxes = torch.cat([v['obj_boxes'] for v in targets])
        bbox1 = box_cxcywh_to_xyxy(out_obj_bbox)
        bbox2 = box_cxcywh_to_xyxy(tgt_obj_boxes)
        cost_obj_giou = -generalized_box_iou(bbox1, bbox2) 
        cost_obj_giou += cost_sub_giou * (tgt_obj_boxes == 0).all(dim=1).unsqueeze(0)

        if cost_sub_giou.shape[1] == 0:
            cost_giou = cost_sub_giou
        else:
            cost_giou = torch.stack((cost_sub_giou, cost_obj_giou)).max(dim=0)[0]
        #   100N, tgtN
        C = self.cost_giou * cost_giou 
        return C       

    def get_cost_matching(self, pred_matching_logits, targets, alpha=0.25, gamma=2, neg_inds_beta=0, hard_neg_inds_beta=3):
        #   100N, 1
        out_matching_prob = pred_matching_logits.flatten(0, 1).softmax(-1)
        tgt_matching_labels = torch.cat([v['matching_labels'] for v in targets])
        if self.loss_type=='ce_loss':
            cost_matching = -out_matching_prob[:, tgt_matching_labels]#1-out_obj_prob[:, tgt_obj_labels]
        elif self.loss_type=='focal_loss':
            neg_cost = (1 - alpha) * (out_matching_prob ** gamma) * (-(1 - out_matching_prob + 1e-8).log())
            pos_cost = alpha * ((1 - out_matching_prob) ** gamma) * (-(out_matching_prob + 1e-8).log())
            cost_matching = pos_cost[:, tgt_matching_labels] - neg_cost[:, tgt_matching_labels]            

        #   100N, tgtN
        C = self.cost_matching * cost_matching
        return C